import warnings
warnings.filterwarnings("ignore")
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"]="3"

import tensorflow as tf
from nets.swin_seg import swin_seg_model
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np


def normalize(img):    
    max_val = np.max(img)
    min_val = np.min(img)
    val_range = max_val - min_val
    norm_0_1 = (img-min_val)/val_range
    img = np.clip(2*norm_0_1-1,-1,1)
    return img



if __name__ == '__main__':
    checkpoint_dir = "logs"
    latest = tf.train.latest_checkpoint(checkpoint_dir)
    # print(latest)
    input_shape = [224,224,3]
    num_classes = 2
    model = swin_seg_model(input_shape=input_shape,
                            num_classes=num_classes)

    model.load_weights(latest).expect_partial()
    print('载入权重成功'.center(60,'='))

    # -----------------------------------
    # 载入图片
    # -----------------------------------
    image_path = "C:/Users/hblee/Documents/datasets/seg2d/train/JPEGImages/ISIC_0000017.jpg"
    label_path = "C:/Users/hblee/Documents/datasets/seg2d/train/Segmentations/ISIC_0000017.png"
    image = Image.open(image_path)    
    label = Image.open(label_path)

    image = image.resize(input_shape[:-1],Image.BICUBIC)
    label = label.resize(input_shape[:-1],Image.NEAREST)

    # 图像转numpy
    image_array = np.array(image)
    # 图像归一化
    image_array = normalize(image_array)
    # 扩展batch维度->[1,224,224,3]
    pred_array = np.expand_dims(image_array,axis=0)

    # -----------------------------------
    # 预测
    # -----------------------------------
    pred_mask = model.predict(pred_array)
    pred_mask = tf.argmax(pred_mask, axis=-1)[0].numpy()
    

    plt.subplot(1,3,1),plt.imshow(image),plt.axis('off'),plt.title('image')
    plt.subplot(1,3,2),plt.imshow(label,cmap='gray'),plt.axis('off'),plt.title('label')
    plt.subplot(1,3,3),plt.imshow(pred_mask,cmap='gray'),plt.axis('off'),plt.title('pred')
    plt.show()

